#!/usr/bin/env python3
"""
net_agent.py
Capture socket connect events (destination ip:port) per PID using tracepoints.
Requires sudo.
"""
from bcc import BPF
from datetime import datetime, timezone
import json, os

OUT_DIR = "/var/log/os_monitor"
os.makedirs(OUT_DIR, exist_ok=True)
OUT_FILE = os.path.join(OUT_DIR, "net.jsonl")

BPF_PROGRAM = r"""
#include <uapi/linux/ptrace.h>
#include <linux/sched.h>
#include <net/sock.h>
#include <bcc/proto.h>

struct data_t {
    u32 pid;
    u64 ts_ns;
    int family;
    u16 dport;
    u32 daddr;
};

BPF_PERF_OUTPUT(events);

int trace_connect(struct pt_regs *ctx, struct sock *sk) {
    struct data_t data = {};
    u16 dport = 0;
    u32 daddr = 0;
    bpf_probe_read_kernel(&dport, sizeof(dport), &sk->__sk_common.skc_dport);
    bpf_probe_read_kernel(&daddr, sizeof(daddr), &sk->__sk_common.skc_daddr);
    data.dport = ntohs(dport);
    data.daddr = daddr;
    data.pid = bpf_get_current_pid_tgid() >> 32;
    data.ts_ns = bpf_ktime_get_ns();
    data.family = sk->__sk_common.skc_family;
    events.perf_submit(ctx, &data, sizeof(data));
    return 0;
}
"""

b = BPF(text=BPF_PROGRAM)
# attach kprobe to tcp_v4_connect or inet_sock_connect
try:
    b.attach_kprobe(event="tcp_v4_connect", fn_name="trace_connect")
except Exception:
    try:
        b.attach_kprobe(event="inet_sock_connect", fn_name="trace_connect")
    except Exception:
        pass


def write_record(record: dict):
    with open(OUT_FILE, "a") as f:
        f.write(json.dumps(record, default=str) + "\n")


def print_event(cpu, data, size):
    event = b["events"].event(data)
    r = {
        "pid": int(event.pid),
        "ts_ns": int(event.ts_ns),
        "family": int(event.family),
        "dport": int(event.dport),
        "daddr": int(event.daddr),
        "daddr_str": __import__('socket').inet_ntoa(int(event.daddr).to_bytes(4, 'big')),
        "ts": datetime.now(timezone.utc).isoformat()
    }
    write_record(r)

b["events"].open_perf_buffer(print_event)
print("[+] net_agent running, writing to {}".format(OUT_FILE))
while True:
    try:
        b.perf_buffer_poll()
    except KeyboardInterrupt:
        break